import numpy as np
from scipy.stats import norm
import matplotlib.pyplot as plt

def importance_sampling_pfa_pd(mu_signal, sigma, threshold, num_samples=10000, shift=1.0):
    """
    使用重要性采样计算虚警概率(P_fa)和检测概率(P_d)
    
    参数:
    mu_signal : float - H1假设下的信号均值
    sigma : float - 噪声标准差
    threshold : float - 检测门限
    num_samples : int - 采样数量
    shift : float - 建议分布的均值偏移量
    
    返回:
    P_fa_est : float - 估计的虚警概率
    P_d_est : float - 估计的检测概率
    P_fa_true : float - 理论虚警概率
    P_d_true : float - 理论检测概率
    """
    # 理论真实值（用于验证）
    P_fa_true = 1 - norm.cdf(threshold, loc=0, scale=sigma)
    P_d_true = 1 - norm.cdf(threshold, loc=mu_signal, scale=sigma)
    
    # 计算虚警概率 P_fa = P(T > threshold | H0)
    # H0建议分布: N(threshold + shift, sigma^2)
    samples_H0 = np.random.normal(loc=threshold + shift, scale=sigma, size=num_samples)
    
    # 计算权重: w = f_H0(x) / q_H0(x)
    weights_H0 = norm.pdf(samples_H0, loc=0, scale=sigma) / \
                 norm.pdf(samples_H0, loc=threshold + shift, scale=sigma)
    
    # 指示函数: 统计量超过门限
    indicator_H0 = (samples_H0 > threshold).astype(float)
    P_fa_est = np.mean(indicator_H0 * weights_H0)
    
    # 计算检测概率 P_d = P(T > threshold | H1)
    # H1建议分布: N(threshold + shift, sigma^2)
    samples_H1 = np.random.normal(loc=threshold + shift, scale=sigma, size=num_samples)
    
    # 计算权重: w = f_H1(x) / q_H1(x)
    weights_H1 = norm.pdf(samples_H1, loc=mu_signal, scale=sigma) / \
                 norm.pdf(samples_H1, loc=threshold + shift, scale=sigma)
    
    indicator_H1 = (samples_H1 > threshold).astype(float)
    P_d_est = np.mean(indicator_H1 * weights_H1)
    
    return P_fa_est, P_d_est, P_fa_true, P_d_true

# 示例参数设置
mu_signal = 2.0    # H1假设下信号均值
sigma = 1.0        # 噪声标准差
threshold = 3.0    # 检测门限
num_samples = 50000 # 重要性采样样本数
shift = 1.0        # 建议分布偏移量

# 计算概率
results = importance_sampling_pfa_pd(mu_signal, sigma, threshold, num_samples, shift)
P_fa_est, P_d_est, P_fa_true, P_d_true = results

# 输出结果
print(f"理论虚警概率 P_fa: {P_fa_true:.6f}")
print(f"估计虚警概率 P_fa: {P_fa_est:.6f}")
print(f"相对误差: {abs(P_fa_est - P_fa_true)/P_fa_true*100:.2f}%")
print("\n" + "-"*50 + "\n")
print(f"理论检测概率 P_d: {P_d_true:.6f}")
print(f"估计检测概率 P_d: {P_d_est:.6f}")
print(f"相对误差: {abs(P_d_est - P_d_true)/P_d_true*100:.2f}%")

# 绘制结果对比
labels = ['理论虚警概率', '估计虚警概率', '理论检测概率', '估计检测概率']
values = [P_fa_true, P_fa_est, P_d_true, P_d_est]

plt.figure(figsize=(10, 6))
plt.bar(labels, values, color=['blue', 'cyan', 'red', 'pink'])
plt.ylabel('概率值')
plt.title('理论值与重要性采样估计值对比')
plt.yscale('log')  # 对数坐标便于观察小概率值
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.show()